package com.kongxiang.cubebit.vectordb.services.impl;

import com.kongxiang.cubebit.vectordb.services.VectorStoreService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;

import java.util.Collection;
import java.util.List;
import java.util.stream.Stream;

@Service
@RequiredArgsConstructor
public class VectorStoreServiceImpl implements VectorStoreService {

    private final VectorStore vectorStore;

    @Override
    public void add(String content) {
        vectorStore.add(List.of(Document.builder().text(content).build()));
    }

    @Override
    public void batchAdd(Collection<String> contents) {
        List<Document> collect = contents.stream()
                .map(s -> Document.builder().text(s).build()).toList();
        vectorStore.add(collect);
    }


    @Override
    public List<Document> search(String query, int k, double threshold) {
        SearchRequest.Builder builder = SearchRequest.builder().query(query);
        if (k > 0) {
            builder.topK(k);
        }
        if (threshold > 0) {
            builder.similarityThreshold(threshold);
        }
        /*if(expression != null){
            builder.filterExpression();
        }*/

        return vectorStore.similaritySearch(builder.build());
    }

    @Override
    public void delete(String id) {
        vectorStore.delete(id);
    }

    @Override
    public void update(String id, String content) {
        vectorStore.accept(List.of(Document.builder().id(id).text(content).build()));
    }
}
